import os,sys
from pylab import *
from scipy.integrate import quad
from scipy.integrate import romberg
from scipy.interpolate import interp1d

from SNeCosmology import SNeCosmology

class PhotozSN:
    def __init__(self,  SN_want=10000, \
                        sigma_int=0.13, \
                        zp_min=0.15, zp_max=1.1, z_sn_max=2.0, \
                        n_mock=100, out_dir=None, \
                        H0=67.7, Omega_m=0.307, \
                        Omega_k=0.0, w0=-1.0, wa=0.0, grid_size=100,\
                        sn_density=None ):
        self.SC = SNeCosmology( H0=H0, Omega_m=Omega_m, Omega_k=Omega_k, \
                                w0=w0, wa=wa, z_sn_max=z_sn_max, \
                                grid_size=grid_size )
                                    
        self.SN_want    = SN_want
        self.n_mock     = n_mock
        self.zp_min     = zp_min
        self.zp_max     = zp_max
        self.sigma_int  = sigma_int

        self.sn_density_file = None
        self.sn_density = None
        
        if sn_density is not None:
            self.sn_density_file = sn_density
            self.sn_density = loadtxt(self.sn_density_file)
        else:
            print("sn_density = None")
            sys.exit(0)

        self.out_dir    = None
        if out_dir is not None:
            self.out_dir= out_dir
            if os.path.isdir(self.out_dir) is False:
                print("out_dir: %s does not exist, so I will create it for you"%self.out_dir)
                os.mkdir(self.out_dir)
        else:
            print("out_dir = None")
            sys.exit(0)
            
        self.print_info()

    def print_info(self):
        print("==> Cosmological Parameters:")
        print("H0       = %g"%self.SC.H0)
        print("Omega_m  = %g"%self.SC.Omega_m)
        print("Omega_k  = %g"%self.SC.Omega_k)
        print("w0       = %g"%self.SC.w0)
        print("wa       = %g"%self.SC.wa)
        print("==> Mock Sample settings:")
        print("SNe number wanted: %d"%self.SN_want)
        print("SNe intrinsic dispersion: %g"%self.sigma_int)
        print("SNe minmum photo-z redshift: %g"%self.zp_min)
        print("SNe maximum photo-z redshift: %g"%self.zp_max)
        print("SNe mock number: %d"%self.n_mock)
        print("SNe density file: %s"%self.sn_density_file)
        print("SNe mock out dir: %s"%self.out_dir)

    def make_sample(self):
        print("\n==> generating mock sample")
        sn_density_max = self.sn_density[:,1].max()*1.05
        zs = linspace(self.sn_density[:,0].min(), self.sn_density[:,0].max(), 201)
        zs = (zs[1:]+zs[:-1])/2.0
        muzs = zeros(zs.shape)

        for i in range(len(zs)):
            muzs[i] = self.SC.dist_mu(zs[i])

        density_interp = interp1d(self.sn_density[:,0],self.sn_density[:,1],kind='linear')
        muzs_interp = interp1d(zs,muzs,kind='linear')

        # update zs_min and zs_max to avoid numerical error
        zs_min = zs.min()
        zs_max = zs.max()

        #####################################################################
        # generating SNe does not need the matrix A, only ns(zs) is needed.
        #####################################################################

        nm = 0
        while nm < self.n_mock:
            cnt = 0
            sample_zp = []
            sample_muzp = []
            sample_muzperr = []
            while cnt < self.SN_want:
                zs_r = zs_min + rand()*(zs_max-zs_min)
                zp_r = zs_r + 0.02*(1+zs_r)*randn()
                r = rand()*sn_density_max
                if r <= density_interp(zs_r):
                    if zp_r >= self.zp_min and zp_r <= self.zp_max:
                        sample_zp.append(zp_r)
                        muzp_err = self.sigma_int * randn()
                        sample_muzp.append(muzs_interp(zs_r)+muzp_err)
                        sample_muzperr.append(muzp_err)
                        cnt += 1

            sample_zp = array(sample_zp).reshape(self.SN_want,1)
            sample_muzp = array(sample_muzp).reshape(self.SN_want,1)
            sample_muzperr = array(sample_muzperr).reshape(self.SN_want,1)

            savetxt(os.path.join(self.out_dir,'MOCK_'+str(nm+1)+'.txt'), \
                    hstack((sample_zp,sample_muzp,sample_muzperr)), \
                    fmt='%10.6f', delimiter=' ')
            nm += 1
            print '--> generated %3d th mock sample'%nm
            pause(1.0)
